#!/usr/bin/env python3

# Copyright 2020 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import json
import random
import secrets
import logging
import hashlib
import argparse

import config.config as pconfig
import utility.logger as plogger
import utility.hex_utils as hex_utils
import avalon_crypto_utils.signature as signature
import avalon_crypto_utils.crypto_utility as crypto_utils
from database import connector
from error_code.error_status import SignatureStatus
from http_client.http_jrpc_client import HttpJrpcClient
import avalon_enclave_manager.base_enclave_info as enclave_info
from avalon_sdk.work_order.work_order_params import WorkOrderParams
from avalon_enclave_manager.worker_kv_delegate import WorkerKVDelegate

logger = logging.getLogger(__name__)


class WPERequester():
    """
    JRPC requester acting on behalf of WorkOrderProcessorEnclaveManager
    """

    def __init__(self, config):
        """
        Constructor for WPERequester. Initialize the HTTP jrpc client.
        Parameters :
            @param config - dict of config read
        """
        self._uri_client = HttpJrpcClient(
            config["KMEListener"]["kme_listener_url"])
        self._conn_retries = config["KMEListener"]["connection_retry"]
        worker_id = config.get("WorkerConfig")["worker_id"]
        # Calculate sha256 of worker id to get 32 bytes. The TC spec proxy
        # model contracts expect byte32. Then take a hexdigest for hex str.
        worker_id = hashlib.sha256(worker_id.encode("UTF-8")).hexdigest()
        if config.get("KvStorage") is None:
            logger.error("Kv Storage path is missing")
            sys.exit(-1)
        try:
            kv_helper = connector.open(config['KvStorage']['remote_url'])
        except Exception as err:
            logger.error("Failed to open KV storage interface; " +
                         "exiting Intel SGX Enclave manager: {}".format(err))
            sys.exit(-1)
        worker_kv_delegate = WorkerKVDelegate(kv_helper)
        self._worker = worker_kv_delegate.get_worker_by_id(worker_id)
        jrpc_methods = {}
        jrpc_methods["kme-uid"] = "GetUniqueVerificationKey"
        jrpc_methods["kme-reg"] = "RegisterWorkOrderProcessor"
        jrpc_methods["kme-preprocess"] = "PreProcessWorkOrder"
        # Map too hold workload-id to JRPC method mapping
        self._jrpc_methods = jrpc_methods

    def get_unique_verification_key(self, verification_key_nonce):
        """
        Request wrapper to get a unique id from the KME

        Parameters :
            @param verification_key_nonce - Random nonce generated by this WPE
        Returns :
            @returns result - Result received from the KME which includes the
                              public verification key which is supposed to be
                              included in REPORTDATA by the WPE. None, in case
                              of failure.
        """
        workload_id = "kme-uid"
        in_data = [verification_key_nonce]

        # Create session key and iv to sign work order request
        session_key = crypto_utils.generate_key()
        session_iv = crypto_utils.generate_iv()

        wo_req = self._construct_wo_req(
            in_data, workload_id, self._worker.encryption_key,
            session_key, session_iv)

        response = self._post_and_get_result(wo_req)

        if "result" in response:
            wo_response_json = response["result"]
            if self._verify_res_signature(wo_response_json,
                                          self._worker.verification_key,
                                          wo_req["params"]["requesterNonce"]):
                decrypted_res = crypto_utils.decrypted_response(
                    wo_response_json, session_key, session_iv)
                # Response contains an array of results. In this case, the
                # array has single element and the data field is of interest.
                # The data contains result,verification_key and
                # verification_key_signature delimited by ' '.
                # @TODO : Update to use multiple out_data fields.
                return decrypted_res[0]['data']
            return None
        else:
            logger.error("Could not get a unique id from the KME : {}"
                         .format(response))
            return None

    def register_wo_processor(self, unique_verification_id,
                              encryption_key,
                              proof_data, mr_enclave):
        """
        Request to register this WPE with the KME

        Parameters :
            @param unique_verification_id - Unique verifying key received from
            KME.
            @param encryption_key - encryption key of WPE
            @param proof_data - The IAS attestation report/DCAP quote
            @param mr_enclave - MRENCLAVE for this WPE
        Returns :
            @returns status - The status of the registration.
                              True, for success. None, in case of errors.
        """
        workload_id = "kme-reg"
        registration_params = {
            "unique_id": unique_verification_id,
            "proof_data": proof_data,
            "wpe_encryption_key": encryption_key,
            "mr_enclave": mr_enclave
        }
        in_data = [json.dumps(registration_params)]

        # Create session key and iv to sign work order request
        session_key = crypto_utils.generate_key()
        session_iv = crypto_utils.generate_iv()

        wo_req = self._construct_wo_req(
            in_data, workload_id, self._worker.encryption_key,
            session_key, session_iv)

        response = self._post_and_get_result(wo_req)

        if "result" in response:
            wo_response_json = response["result"]
            if "error" not in wo_response_json and self._verify_res_signature(
                    wo_response_json, self._worker.verification_key,
                    wo_req["params"]["requesterNonce"]):
                decrypted_res = crypto_utils.decrypted_response(
                    wo_response_json, session_key, session_iv)
                # Response contains an array of results. In this case, the
                # array has single element and the data field is of interest.
                # It is integer with status of registration.
                return decrypted_res[0]['data']
            return None
        else:
            logger.error("Could not register this WPE with the KME : {}"
                         .format(response))
            return None

    def preprocess_work_order(self, wo_request, encryption_key):
        """
        Request to preprocess a work order

        Parameters :
            @param wo_request - The original work order request as str
            @param encryption_key - WPE's public encryption key
        Returns :
            @returns result - Result from KME that includes the workorder
                              key info. error response, in case of failure.
        """
        workload_id = "kme-preprocess"
        in_data = [wo_request, encryption_key]

        # Create session key and iv to sign work order request
        session_key = crypto_utils.generate_key()
        session_iv = crypto_utils.generate_iv()

        wo_req = self._construct_wo_req(
            in_data, workload_id, self._worker.encryption_key,
            session_key, session_iv)

        response = self._post_and_get_result(wo_req)

        if "result" in response:
            wo_response_json = response["result"]
            if self._verify_res_signature(wo_response_json,
                                          self._worker.verification_key,
                                          wo_req["params"]["requesterNonce"]):
                decrypted_res = crypto_utils.decrypted_response(
                    wo_response_json, session_key, session_iv)
                # Response contains an array of results. In this case, the
                # array has single element and the data field is of interest.
                return decrypted_res[0]['data']
            return None
        else:
            logger.error("Could not preprocess work order at KME : {}"
                         .format(response))
            return response

    def _construct_wo_req(self, in_data, workload_id, encryption_key,
                          session_key, session_iv):
        """
        Construct the parameters for a standard work order request

        Parameters :
            @param in_data - List of data to be passed to workload processor
            @param workload_id - Id of the target workload
            @encryption_key - Worker encryption key
            @session_key - Session key to be embedded in request
            @session_iv - Session key iv for encryption algorithm
        Returns :
            @returns A json request prepared using parameters passed
        """
        # Create work order
        # Convert workloadId to hex
        workload_id_hex = workload_id.encode("UTF-8").hex()
        work_order_id = secrets.token_hex(32)
        requester_id = secrets.token_hex(32)
        requester_nonce = secrets.token_hex(16)
        # worker id is not known here. Hence passing a random string
        worker_id = secrets.token_hex(32)
        # Create work order params
        wo_params = WorkOrderParams()
        wo_params.create_request(
            work_order_id, worker_id, workload_id_hex, requester_id,
            session_key, session_iv, requester_nonce,
            result_uri=" ", notify_uri=" ",
            worker_encryption_key=encryption_key,
            data_encryption_algorithm="AES-GCM-256"
        )
        for value in in_data:
            wo_params.add_in_data(value)

        # Encrypt work order request hash
        wo_params.add_encrypted_request_hash()

        return {
            "jsonrpc": "2.0",
            "id": random.randint(0, 100000),
            "params": json.loads(wo_params.to_string()),
            "method": self._jrpc_methods[workload_id]
        }

    def _verify_res_signature(self, work_order_res,
                              worker_verification_key,
                              requester_nonce):
        """
        Verify work order result signature

        Parameters :
            @param work_order_res - Result from work order response
            @param worker_verification_key - Worker verification key
            @param requester_nonce - requester generated nonce in work
            order request
        Returns :
            @returns True - If verification succeeds
                    False - If verification fails
        """
        sig_obj = signature.ClientSignature()
        status = sig_obj.verify_signature(
            work_order_res,
            worker_verification_key,
            requester_nonce)
        if status == SignatureStatus.PASSED:
            logger.info("Signature verification successful")
        else:
            logger.error("Signature verification failed")
            return False
        return True

    def _post_and_get_result(self, json_rpc_request):
        """
        Helper method to serialize and send JRPC request and get response
        from the KME.

        Parameters :
            @param json_rpc_request - JSON containing RPC request
        Returns :
            @returns response - Response received from the KME
        """
        json_request_str = json.dumps(json_rpc_request)
        logger.info("Request to KME listener %s", json_request_str)
        try:
            # Attempt to post request could fail if KME yet to be ready to
            # receive requests. Hence pass in the number of retries(>1)
            # to post a request.
            response = self._uri_client._postmsg(json_request_str,
                                                 self._conn_retries)
        except Exception as err:
            logger.error("Exception occurred in communication with KME")
            raise err
        logger.info("Response from KME %s", response)

        return response
